
#export plots and tables
function pad(series,inds)
    if lastindex(series) >= lastindex(inds)
        out = series[inds]
    else
        out = [series; [series[end] for t=length(series)+1:length(inds)]]
    end
    out
end

nIter_plots = 30
tab_means = c_means[1][:,[:inds,]]
tab_std = c_means[1][:,[:inds,]]
tab_std_dif = c_means[1][:,[:inds,]]
byIter_means = DataFrame(inds=1:nIter_plots)
byIter_std = DataFrame(inds=1:nIter_plots)
byIter_std_dif = DataFrame(inds=1:nIter_plots)

for k in names(tab_means)
    if occursin("graduation",k) || occursin("enroll",k)
        tab_means[!,k] .*= 100
        tab_std[!,k] .*= 100
        tab_std_dif[!,k] .*= 100
    end
end

for (col,strkey) in enumerate(names(c_means[1]))
    col==1 && continue
    key = Symbol(strkey)
    iswelfare = occursin("welfare",strkey)
    isenroll = occursin("enroll",strkey)
    isgrad = occursin("graduation",strkey)
    @assert iswelfare || isenroll || isgrad
    #referenceSeries = iswelfare ? :benchmark_welfare : :benchmark_enroll
    if iswelfare
        referenceSeries = :efficient_welfare
    elseif isenroll
        referenceSeries = :efficient_enroll
    else
        referenceSeries = :efficient_graduation
    end
    # create mean and std of estimates
    vals = hcat([c_means[t][!,col] for t=1:length(c_means)]...)
    vals_dif = hcat([c_means[t][!,col] .- c_means[t][:,referenceSeries] for t=1:length(c_means)]...)
    tab_means[!,key] = vec(mean(vals,dims=2))
    tab_std[!,key] = vec(std(vals,dims=2))
    tab_std_dif[!,key] = vec(std(vals_dif,dims=2))
    if !iswelfare
        tab_means[!,key] .*= 100
        tab_std[!,key] .*= 100
        tab_std_dif[!,key] .*= 100
    end
    #
    if !isgrad
        inds = byIter_means[!,:inds]
        vals = hcat([pad(c_byIter[t][key],inds) for t=1:length(c_byIter)]...)
        vals_dif = hcat([pad(c_byIter[t][key],inds) .- pad(c_byIter[t][referenceSeries],inds) for t=1:length(c_byIter)]...)
        byIter_means[!,key] = vec(mean(vals,dims=2))
        byIter_std[!,key] = vec(std(vals,dims=2))
        byIter_std_dif[!,key] = vec(std(vals_dif,dims=2))
        if !iswelfare
            byIter_means[!,key] .*= 100
            byIter_std[!,key] .*= 100
            byIter_std_dif[!,key] .*= 100
        end
    end
end

function makeByIterWelfarePlot()
    plt = plot()
    nIter = [length(c[:benchmark_enroll]) for c in c_byIter]
    plot!(plt,byIter_means.inds, byIter_means.benchmark_welfare,
        ribbon=1.96 .* byIter_std_dif.benchmark_welfare,label="2012 Baseline: $(Int(round(mean(nIter)))) iterations")
    plot!(plt,byIter_means.inds,byIter_means.smallalpha2012_welfare,
        ribbon=1.96 .* byIter_std_dif.smallalpha2012_welfare,label="alpha=.1")
    plot!(plt,byIter_means.inds,byIter_means.zeroalpha_welfare,
        ribbon=1.96 .* byIter_std_dif.zeroalpha_welfare,label="alpha=0")
    plot!(plt,byIter_means.inds,byIter_means.asif2011_welfare,
        ribbon=1.96 .* byIter_std_dif.asif2011_welfare,label="Exclude G8")
    plot!(plt,byIter_means.inds,byIter_means.smallalpha_welfare,
        ribbon=1.96 .* byIter_std_dif.smallalpha_welfare,label="Exclude G8, alpha=.1")
    # plot!(plt,byIter_means.inds,byIter_means.smallalpha_welfare,
    #     ribbon=1.96 .* byIter_std_dif.equalalpha_welfare,label="Exclude G8, equal alpha")
    plot!(plt,byIter_means.inds,byIter_means.efficient_welfare,label="Efficient bound",
        linestyle=:dash, legend=:bottomright)
    #=
    plot!(plt,byIter_means.inds,byIter_means.asInData_welfare,
        ribbon=1.96 .* byIter_std_dif.asInData_welfare, label="Enrollment as Observed",
        linestyle=:dash, legend=:bottomright)
    nIter_observed = findfirst(byIter_means.benchmark_welfare .> byIter_means.asInData_welfare)
    plot!(plt,[nIter_observed],seriestype="vline",linestyle=:dashdot, label="\"observed\" iterations",
        legend=:bottomright)
    =#
    xlabel!(plt,"aftermarket DA iteration")
    ylabel!(plt,"mean welfare (1m Chilean Pesos)")
    savefig(plt,figurepath*"/welfareByIter.png")
end
makeByIterWelfarePlot()


function makeByIterEnrollPlot()
    plt = plot()
    nIter = [length(c[:benchmark_enroll]) for c in c_byIter]
    plot!(plt,byIter_means.inds, byIter_means.benchmark_enroll,
        ribbon=1.96 .* byIter_std_dif.benchmark_enroll,label="2012 Baseline: $(Int(round(mean(nIter)))) iterations")
    plot!(plt,byIter_means.inds,byIter_means.smallalpha2012_enroll,
        ribbon=1.96 .* byIter_std_dif.smallalpha2012_enroll,label="alpha=.1")
    plot!(plt,byIter_means.inds,byIter_means.zeroalpha_enroll,
        ribbon=1.96 .* byIter_std_dif.zeroalpha_enroll,label="alpha=0")
    plot!(plt,byIter_means.inds,byIter_means.asif2011_enroll,
        ribbon=1.96 .* byIter_std_dif.asif2011_enroll,label="Exclude G8")
    plot!(plt,byIter_means.inds,byIter_means.smallalpha_enroll,
        ribbon=1.96 .* byIter_std_dif.smallalpha_enroll,label="Exclude G8, alpha=.1")
    plot!(plt,byIter_means.inds,byIter_means.efficient_enroll,label="Efficient bound",
    linestyle=:dash, legend=:bottomright)
    # plot!(plt,byIter_means.inds,byIter_means.asInData_enroll,
    #     ribbon=1.96 .* byIter_std_dif.asInData_enroll, label="Enrollment as Observed", linestyle=:dash)
    # nIter_observed = findfirst(byIter_means.benchmark_enroll .> byIter_means.asInData_enroll)
    # plot!(plt,[nIter_observed],seriestype="vline",linestyle=:dashdot, label="\"observed\" iterations",
    #     legend=:bottomright)
    xlabel!(plt,"aftermarket DA iteration")
    ylabel!(plt,"Percent enrolled")
    savefig(plt,figurepath*"/EnrollByIter.png")
end
makeByIterEnrollPlot()



function makeWelfareTable(f=stdout)
    println(f,"\\begin{tabular}{lccccc}")
    println(f,"\\hline\\hline")
    println(f,"Counterfactual & All (Avg.) & Male Private & Male Public & Female Private & Female Public \\\\ \\hline")
    println(f,"& & & & & \\\\")
    #
    tm = tab_means[!,:efficient_welfare]
    println(f,"A. Welfare & & & & & \\\\")
    minirow("No Frictions (*)",tm,tab_std[!,:efficient_welfare],f)
    welfareTableRow("Baseline - *",:benchmark_welfare,tm,f)
    # welfareTableRow("\$\\alpha=.1\$ - Bound",:smallalpha2012_welfare,tm,f)
    # welfareTableRow("\$\\alpha=0\$ - Bound",:zeroalpha_welfare,tm,f)
    welfareTableRow("Exclude G8 - *",:asif2011_welfare,tm,f)
    # welfareTableRow("Exclude G8, \$\\alpha=.1\$ - Bound",:smallalpha_welfare,tm,f)
    welfareTableRow("Exclude G8, Equal \$\\alpha\$ - *",:equalalpha_welfare,tm,f)
    # welfareTableRow("Free college for low SES - Bound",:freecollege_lowSES_welfare,tm,f)
    # welfareTableRow("Free college for low SES, Exclude G8 - Bound",:freecollege_lowSES_asif2011_welfare,tm,f)
    #welfareTableRow("Observed - Bound",:asInData_welfare,tm,f)
    #
    tm = tab_means[!,:efficient_enroll]
    println(f,"B. Enrollment (pct) & & & & & \\\\")
    minirow("No Frictions (*)",tm,tab_std[!,:efficient_enroll],f)
    welfareTableRow("Baseline - *",:benchmark_enroll,tm,f)
    # welfareTableRow("\$\\alpha=.1\$ - Bound",:smallalpha2012_enroll,tm,f)
    # welfareTableRow("\$\\alpha=0\$ - Bound",:zeroalpha_enroll,tm,f)
    welfareTableRow("Exclude G8 - *",:asif2011_enroll,tm,f)
    # welfareTableRow("Exclude G8, \$\\alpha=.1\$ - Bound",:smallalpha_enroll,tm,f)
    welfareTableRow("Exclude G8, Equal \$\\alpha\$ - *",:equalalpha_enroll,tm,f)
    # welfareTableRow("Free college for low SES - Bound",:freecollege_lowSES_enroll,tm,f)
    # welfareTableRow("Free college for low SES, Exclude G8 - Bound",:freecollege_lowSES_asif2011_enroll,tm,f)
    #welfareTableRow("Observed - Bound",:asInData_enroll,tm,f)
    #
    tm = tab_means[!,:efficient_graduation]
    println(f,"C. Six-Year Graduation (pct) & & & & & \\\\")
    minirow("No Frictions (*)",tm,tab_std[!,:efficient_graduation],f)
    welfareTableRow("Baseline - *",:benchmark_graduation,tm,f)
    # welfareTableRow("\$\\alpha=.1\$ - Bound",:smallalpha2012_graduation,tm,f)
    # welfareTableRow("\$\\alpha=0\$ - Bound",:zeroalpha_graduation,tm,f)
    welfareTableRow("Exclude G8 - *",:asif2011_graduation,tm,f)
    # welfareTableRow("Exclude G8, \$\\alpha=.1\$ - Bound",:smallalpha_graduation,tm,f)
    welfareTableRow("Exclude G8, Equal \$\\alpha\$ - *",:equalalpha_graduation,tm,f)
    # welfareTableRow("Free college for low SES - Bound",:freecollege_lowSES_graduation,tm,f)
    # welfareTableRow("Free college for low SES, Exclude G8 - Bound",:freecollege_lowSES_asif2011_graduation,tm,f)
    #welfareTableRow("Observed - Bound",:asInData_graduation,tm,f)
    #
    println(f," & & & & & \\\\ \\hline \\hline")
    println(f,"\\end{tabular}")
end

welfareTableRow(name,series,tm,f) = minirow(name,tab_means[!,series].-tm,tab_std_dif[!,series],f)

open(tablepath*"/mainresults.tex","w") do f
    makeWelfareTable(f)
end

function makeExtraWelfareTable(f,names,vars)
    println(f,"\\begin{tabular}{lccccc}")
    println(f,"\\hline\\hline")
    println(f,"Counterfactual & All (Avg.) & Male Private & Male Public & Female Private & Female Public \\\\ \\hline")
    #
    for (letter,outcome) in zip(
            ["A. Welfare (1m CLP)", "B. Enrollment (pct)", "C. Six-year Graduation (pct)"],
            ["welfare", "enroll", "graduation"])
        println(f,"& & & & & \\\\")
        tm = tab_means[!,Symbol("efficient_$outcome")]
        println(f,"$(letter) & & & & & \\\\")
        minirow("No Frictions (*)",tm,tab_std[!,Symbol("efficient_$outcome")],f)
        for (k,v) in zip(names,vars)
            welfareTableRow("$k - *",Symbol("$(v)_$(outcome)"),tm,f)
        end
    end
    println(f," & & & & & \\\\ \\hline \\hline")
    println(f,"\\end{tabular}")
end

open(tablepath*"/additionalresults.tex","w") do f
    makeExtraWelfareTable(f,["Free G33", "Free G25, G8 off", "Max Welfare"],
        [:freecollege_lowSES,:freecollege_lowSES_asif2011, :maxutility])
end

open(tablepath*"/newmainresults.tex","w") do f
    makeExtraWelfareTable(f,["Baseline", "Exclude G8", "Max Welfare", "Free G33"],
        [:benchmark, :asif2011, :maxutility, :freecollege_lowSES,])
end

function makeAlphaFigures(vartype="welfare",typ=0,xlab="fraction reduction in frictions")
    (cf_constants,cf_temps) = counterfactuals[2012]
    (cc, _, _) = dataset[2012]
    nj = 9
    onp = hcat([tab_means[!,Symbol("scaleAlpha_$(dec)_$vartype")] for dec=0:9]...,
        tab_means[!,Symbol("benchmark_$vartype")])
    offp = hcat([tab_means[!,Symbol("scaleAlpha2011_$(dec)_$vartype")] for dec=0:9]...,
        tab_means[!,Symbol("asif2011_$vartype")])
    plt = scatter(1 .- (0:.1:1),onp[typ+1,:], label = "All On Platform")
    scatter!(plt,1 .- (0:.1:1),offp[typ+1,:], label = "Exclude G8", legend=:bottomright)
    xlabel!(plt,xlab)
    plt
end
for vartype in ["welfare","enroll","graduation"]
    plt = makeAlphaFigures(vartype,0)
    savefig(plt,figurepath*"/shrinkAlpha_$vartype.png")
end

plts = [makeAlphaFigures("welfare",t,"Welfare, $(typenames[t])") for t=1:4 ]
for plt in plts[1:end-1]
    plot!(plt,legend=:false)
end
plt = plot(plts...)
savefig(plt,figurepath*"/shrinkAlpha_welfare_heterogeneity.png")

function makeDrop1Figures(vartype="welfare") #try graduation
    (cf_constants,cf_temps) = counterfactuals[2012]
    (cc, _, _) = dataset[2012]
    nInst = length(unique(cf_constants.institutionID))
    welf_drop = hcat([tab_means[!,Symbol("dropInstitution$(jj)_$vartype")] for jj=1:nInst]...)
    selectivity = zeros(nInst)
    for jj=1:nInst
        num = 0.0
        denom = 0
        programs_j = findall(cf_constants.institutionID .== jj)
        for k in programs_j
            num += sum(cf_constants.priorityEligible[(cc.enroll .== k) .& .!(cc.specialAdmit), k])
            denom += sum(.!(cc.specialAdmit) .& (cc.enroll .== k))
        end
        selectivity[jj] = num/denom
    end
    inds = sortperm(selectivity)
    plt_pooled = scatter(1:size(welf_drop,2), welf_drop[1,inds] .-  tab_means[1,Symbol("benchmark_$vartype")], label="$vartype, drop institution j",legend=:topleft)
    plts = Any[]
    for typ=1:4
        plt = scatter(1:size(welf_drop,2),welf_drop[typ+1,inds] .- tab_means[typ+1,Symbol("benchmark_$vartype")],legend=:false)
        title!(plt,typenames[typ])
        push!(plts,plt)
    end
    return selectivity, plt_pooled, plts
end

function makeDropDecileTables()
    (cf_constants,cf_temps) = counterfactuals[2012]
    (cc, _, _) = dataset[2012]
    nj = length(unique(programselectivitybin))
    df_out = DataFrame(type=["Pooled",1,2,3,4])
    for vartype in ["welfare","enroll","graduation"]
        welf_drop = hcat([tab_means[!,Symbol("dropSelectivityBin$(dec)_$vartype")] for dec=1:nj]...)
        dif = welf_drop .- tab_means[!,Symbol("benchmark_$vartype")]
        df_out[!,Symbol("benchmark_$vartype")] = tab_means[!,Symbol("benchmark_$vartype")]
        for j=1:nj
            df_out[!,Symbol("dropdecile$(j)_$vartype")] = welf_drop[:,j]
        end
    end
    df_out
end



function makeDropDecileFigures(vartype="welfare") #try graduation
    error("don't run me")
    (cf_constants,cf_temps) = counterfactuals[2012]
    (cc, _, _) = dataset[2012]
    nj = length(unique(programselectivitybin))
    welf_drop = hcat([tab_means[!,Symbol("dropSelectivityBin$(dec)_$vartype")] for dec=1:nj]...)
    selectivity = zeros(nj)
    for jj=1:nj
        num = 0.0
        denom = 0
        programs_j = findall(cf_constants.institutionID .== jj)
        for k in programs_j
            num += sum(cf_constants.priorityEligible[(cc.enroll .== k) .& .!(cc.specialAdmit), k])
            denom += sum(.!(cc.specialAdmit) .& (cc.enroll .== k))
        end
        selectivity[jj] = num/denom
    end
    inds = sortperm(selectivity)
    plt_pooled = scatter(1:size(welf_drop,2), welf_drop[1,inds] .-  tab_means[1,Symbol("benchmark_$vartype")],legend=:false)
    xlabel!(plt_pooled,"Decile of selectivity")
    title!(plt_pooled,"$vartype: drop programs in selectivity decile - 2012 baseline")
    plts = Any[]
    for typ=1:4
        plt = scatter(1:size(welf_drop,2),welf_drop[typ+1,inds] .- tab_means[typ+1,Symbol("benchmark_$vartype")],legend=:false)
        title!(plt,typenames[typ])
        push!(plts,plt)
    end
    return selectivity, plt_pooled, plts
end

# for vartype in ["welfare","enroll","graduation"]
#     selectivity, plt_drop1_byselectivity, plts = makeDrop1Figures(vartype)
#     plt_drop1 = plot(plts...,layout=(2,2)) #title!(plt_drop1,"Welfare loss: Drop one institution - Baseline")
#     savefig(plt_drop1,figurepath*"/DropOneInstitution_$(vartype)_byType.png")
#     savefig(plt_drop1_byselectivity,figurepath*"/DropOneInstitution_$(vartype).png")
# end


# for vartype in ["welfare","enroll","graduation"]
#     selectivity, plt_drop1_byselectivity, plts = makeDropDecileFigures(vartype)
#     plt_drop1 = plot(plts...,layout=(2,2)) #title!(plt_drop1,"Welfare loss: Drop one institution - Baseline")
#     savefig(plt_drop1,figurepath*"/DropOneDecile_$(vartype)_byType.png")
#     savefig(plt_drop1_byselectivity,figurepath*"/DropOneDecile_$(vartype).png")
# end

# df_dropdecile = makeDropDecileTables()
CSV.write(outputpath*"/counterfactuals_means_bytype.csv",tab_means)
CSV.write(outputpath*"/counterfactuals_sd_bytype.csv",tab_std)
# CSV.write(outputpath*"/counterfactuals_sd_dif_bytype.csv",tab_std_dif)
#


# let
#     cf_constants, cf_temps = counterfactuals[2012]
#     (cc,lv,_) = dataset[2012]
#     J = size(cc.J_lb,1)
#     selectivity = [ mean(cf_constants.priorityEligible[ (cc.enroll .== jj) .& .!cc.specialAdmit, jj])  for jj=1:J]
# end
